Conversation
github-actions
Bot
added
bug
andygrove
left a comment
andygrove
left a comment
There was a problem hiding this comment.
Thanks for this. I think it is the right fix for #5961. Since DataFusion 53 (90633dc), try_grow compares the pool-wide state.used against pool_size / num, so every newly registered consumer tightens the ceiling on the ones already running. This PR goes back to a per-consumer share, charges siblings from new_empty/split/take against that share, and adds the separate pool-total cap that was asked for on #5466. The tests go through the real admission path, and the 10+6 then +10 case fails on main, which is exactly the regression. Could you add "Closes #5961"? I'd plan to close #5466 as superseded once this lands.
I checked a few things that look correct: consumer ids are unique per process and MemoryConsumer is not Clone in DataFusion 55.1, so the duplicate-register assert can't fire; unregister always follows the final free; and a failed or partial grant leaves both counters untouched.
Some things before this leaves draft:
1. Ordering with #5613. This PR, #5466 and #5613 all rewrite register/unregister/shrink/try_grow in fair_pool.rs and conflict with each other. I've requested changes on #5613 and would like to land it first, then rebase this onto it. After the rebase:
consumer_usageneeds to be charged in #5613's locked booking step. It also needs rolling back on every failure path there (finish_acquire(0, additional)on JVM error, short grant and panic). If a path is missed,consumer_usageleaks and item 2 below turns into a panic when the consumer unregisters.- The
#[cfg(test)] test_memory_managerfield (fair_pool.rs:41) should use #5613'sTaskMemoryBridgestub instead, so there is one test seam.
2. A failed release can now abort the executor (fair_pool.rs:135-162).
shrinkpanics on a failed JNI release (:157) beforeconsumer_usageis decremented (:158).- When that happens inside
MemoryReservation::dropand the reservation holds the lastSharedRegistration,unregisterruns during the unwind. There,assert_eq!(usage, Some(0))(:137) panics again, which is a panic during a panic, so the process aborts. - On main the same failure is a single panic that gets converted to a Java exception.
- Decrementing before the release (which is what #5613 does) avoids this, or
unregistercoulddebug_assertand log a warning instead.
3. Evidence. This changes when every fair_unified query spills, so it needs a real off-heap run. Compared with main, the new check never rejects anything main accepts: whenever used + additional <= pool_size / num passes, both new conditions pass too. So it can only reduce spills and let tasks reserve more. Since April, deployments may have tuned memoryPool.fraction against the tighter accidental cap. Compared with per-reservation checking (#5466, pre-DF53, upstream FairSpillPool), it is stricter for DataFusion sort, merge and aggregate, which use sibling reservations. A multi-core off-heap TPC-H run comparing spill counts and RSS (or jemalloc_allocated against summed reservations) against main would settle it. This is probably also worth a line in the release notes.
4. Docs. Several places describe the current pool-wide comparison and would go stale:
docs/source/contributor-guide/memory_management.md(the fair pool description with the 8 GiB / 3 GiB / 2 GiB example).ai/skills/review-comet-memory-pr/SKILL.md(thefair_unifiedchecklist item)docs/source/user-guide/latest/tuning.md(thefair_unifieddescription, which also saysnum_reservationswhere the divisor is registered consumers)
They should describe the new rules: a per-consumer share shared across sibling reservations, a separate pool-total cap, and how this differs from upstream FairSpillPool.
5. Small note on the pool-total cap (non-blocking).
pool_sizeis the executor-wideoffHeap.size * memoryPool.fraction, but each task gets its own pool of that size.- At the default
fraction = 1.0, Spark's per-task cap is always tighter, so the new local cap only binds withfraction < 1and fewer than1/fractionactive tasks. - That is fine, but the comment at
:186-187and the docs shouldn't imply it enforces an executor-wide limit.
|
A correction to item 2 in my earlier review. #5613 changed shape in |
2 participants
Which issue does this PR close?
Addresses the fair-memory-pool accounting in #5212. Overlaps #5466 and needs merge-order coordination with #5613.
Rationale for this change
The fair pool can force operators to spill even when they have room within their own allowance. Consider a 32 MiB pool shared by two consumers, each entitled to 16 MiB. If one holds 10 MiB and the other holds 6 MiB, the second should be able to acquire another 10 MiB. Today, Comet compares the pool's combined 16 MiB usage against that consumer's 16 MiB allowance and rejects the request. With more consumers, this can leave an increasingly large part of the configured pool unusable.
The correction also needs to distinguish a consumer, which owns the allowance, from its reservations, which account for individual buffers. One operator can hold several reservations under the same consumer registration. Checking only the reservation making the request would still let a consumer whose sibling reservations already hold 10 MiB and 6 MiB acquire more through an empty sibling, despite having exhausted its 16 MiB share.
There is a separate total-capacity constraint. An operator might fill the pool before a second consumer registers. The newcomer has a fair share on paper, but cannot allocate until some existing memory is released.
What changes are included in this PR?
The pool now treats all reservations belonging to one consumer as a single account. Growth is admitted only when it fits both that consumer's fair share and the space remaining in the whole pool. Splitting a reservation, moving its contents, or creating an empty sibling preserves the existing charge; releasing memory restores the consumer's available allowance. The pool also reports its configured finite capacity.
This preserves the current policy of dividing the pool among all registered consumers; changing how spillable consumers participate remains in #5465. This overlaps the fair-share correction in #5466 and extends it to cover siblings and total capacity. #5613 separately changes locking around JVM memory calls. These changes need an agreed merge order so its admission and rollback paths retain the same accounting guarantees.
How are these changes tested?
Seven regression tests exercise the actual pool admission logic with DataFusion reservations and simulated JVM grants. They cover the examples above, reservation splitting and transfer, release and retry, mixed consumers, oversized requests, and failed or partial grants.
The original PR revision passed native Rust CI. Its only failed test was a shared decimal codegen coverage assertion, which is corrected by upstream #5849 and now included in this branch. The assertion was reproduced before the correction and passed afterward in a focused local run; local formatting checks also passed.
See fresh CI for revision
865529d4for validation of the updated branch. A dedicated Spark off-heap integration run has not been performed.